// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#pragma once

#include <list>
#include <map>
#include <memory>
#include <mutex>  // NOLINT
#include <utility>

#include "paddle/phi/core/memory/allocation/allocator.h"
#include "paddle/phi/core/memory/allocation/spin_lock.h"

COMMON_DECLARE_bool(enable_auto_growth_allocator_add_lock);

namespace paddle {
namespace memory {
namespace allocation {

class PADDLE_API AutoGrowthBestFitAllocator : public Allocator {
 public:
  AutoGrowthBestFitAllocator(std::shared_ptr<Allocator> underlying_allocator,
                             size_t alignment,
                             size_t chunk_size = 0,
                             bool allow_free_idle_chunk = true,
                             int extra_padding_size = 0);

  bool IsAllocThreadSafe() const override { return true; }

  void DumpInfo() const;

  void PreAlloc() override;

 protected:
  phi::Allocation *AllocateImpl(size_t size) override;

  void FreeImpl(phi::Allocation *allocation) override;

  bool is_small_free_block(size_t size);
  size_t auto_growth_size(bool is_small, size_t chunk_size);

  // Release the memory block which is not used in pool.
  uint64_t ReleaseImpl(const phi::Place &place) override {
    // TODO(vivienfanghuagood): the next line may cause the process to deadlock.
    if (FLAGS_enable_auto_growth_allocator_add_lock) {
      std::lock_guard<SpinLock> guard(spinlock_);
      return FreeIdleChunks();
    }
    return FreeIdleChunks();
  }

 protected:
  uint64_t FreeIdleChunks();
  void Trace() const;

  template <typename T>
  using List = std::list<T>;

  struct Chunk;

  struct Block {
    Block(void *ptr, size_t size, bool is_free, bool is_small, Chunk *chunk)
        : ptr_(ptr),
          size_(size),
          is_free_(is_free),
          is_small_(is_small),
          chunk_(chunk) {}

    void *ptr_;
    size_t size_;
    bool is_free_;
    bool is_small_;
    Chunk *chunk_;  // which chunk it is from
  };

  struct Chunk {
    explicit Chunk(DecoratedAllocationPtr allocation)
        : allocation_(std::move(allocation)) {}

    DecoratedAllocationPtr allocation_;
    List<Block> blocks_;
  };

  struct BlockAllocation : public Allocation {
    explicit BlockAllocation(const List<Block>::iterator &it)
        : Allocation(it->ptr_,
                     it->chunk_->allocation_->base_ptr(),
                     it->size_,
                     it->chunk_->allocation_->place()),
          block_it_(it) {}

    List<Block>::iterator block_it_;
  };

  using BlockIt = List<Block>::iterator;

  std::shared_ptr<Allocator> underlying_allocator_;
  std::map<std::pair<size_t, void *>, BlockIt> small_free_blocks_;
  std::map<std::pair<size_t, void *>, BlockIt> large_free_blocks_;
  std::list<Chunk> chunks_;
  size_t alignment_;
  size_t chunk_size_;
  bool allow_free_idle_chunk_;
  int extra_padding_size_;

  // stat info
  size_t total_alloc_times_;
  size_t total_alloc_size_;
  size_t total_free_times_;
  size_t total_free_size_;

  SpinLock spinlock_;
};

}  // namespace allocation
}  // namespace memory
}  // namespace paddle
